# 读取files目录下的所有文件及第一级目录下文件，进行OCR识别
# 将识别结果输出到output目录的对应目录中， 如files下有合同目录，那就在output下创建对应的合同目录
# 创建识别结果文件 文件名.后缀名.txt 

from paddleocr import PaddleOCR
import os
from colorama import Fore
import time

# 记录程序开始时间
start_time = time.time()

# 初始化 PaddleOCR 模型
ocr = PaddleOCR(use_angle_cls=True, lang="ch")

# 指定输入和输出目录
input_dir = './files'   # 输入目录
output_dir = './output'  # 输出目录

# 确保输出目录存在，如果不存在则创建
os.makedirs(output_dir, exist_ok=True)

# 统计文件的数量
total_files = 0
processed_files = 0
skipped_files = 0

# 遍历文件目录及其子目录
for root, dirs, files in os.walk(input_dir):
    # 获取相对于输入目录的路径
    relative_path = os.path.relpath(root, input_dir)
    # 构造对应的输出目录路径
    output_subdir = os.path.join(output_dir, relative_path)
    # 确保输出子目录存在，如果不存在则创建
    os.makedirs(output_subdir, exist_ok=True)
    
    # 遍历当前目录下的所有文件
    for filename in files:
        # 拼接输入文件的完整路径
        input_file = os.path.join(root, filename)
        # 更新文件总数
        total_files += 1

        # 判断文件扩展名是否为支持 OCR 的格式
        supported_formats = ['.png', '.jpg', '.jpeg', '.pdf']
        ext = os.path.splitext(filename)[1]
        if ext.lower() in supported_formats:
            # 构造输出文件名 文件名.后缀名.txt
            output_filename = filename + ".txt"
            output_file = os.path.join(output_subdir, output_filename)

            # 检查是否已存在相应的文本文件，如果存在说明已经识别过，跳过当前文件
            if os.path.exists(output_file):
                skipped_files += 1
                print(Fore.RED + f"==> 跳过文件 {input_file}: 文件已经识别")
                continue

            # 打印绿色提示
            print(Fore.BLUE +f"==> 开始识别文件 {input_file}")

            # 使用 PaddleOCR 模型进行文字识别
            result = ocr.ocr(input_file, cls=True)

            # 将识别结果拼接成一个字符串
            text_result = ""
            for idx in range(len(result)):
                res = result[idx]
                for line in res:
                    text_result += line[1][0] + "\n"

            # 将识别结果写入到文本文件中
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(text_result)

            print(Fore.GREEN + f"==> 文件 {filename} 识别完成，识别结果已写入 {output_file}")
            # 更新已处理的文件数量
            processed_files += 1
        else:
            # 更新跳过的文件数量
            skipped_files += 1
            print(Fore.RED + f"==> 跳过文件 {filename}: 文件类型不支持 OCR")

# 计算程序执行时间
end_time = time.time()
total_time_seconds = end_time - start_time

# 输出程序执行情况及总耗时
print(Fore.YELLOW + f"==> 文件总数： {total_files} ，识别数：{processed_files} ，跳过：{skipped_files} ")
print(Fore.YELLOW + f"==> 本次运行总耗时：{total_time_seconds:.2f} 秒")





